#include <QMessageBox>
#include <QProcess>
#include <QScrollBar>
#include "classes/chathistory.h"
#include "mainwindow.h"
#include "classes/historystorage.h"
#include "ui_mainwindow.h"
#include "utils/xio.h"

MainWindow::MainWindow(QWidget *parent)
    : QMainWindow(parent)
    , ui(new Ui::MainWindow)
{
    ui->setupUi(this);
    this->loadComboBoxModelPicker(ui->comboBoxModelTypePicker->currentIndex());
}

MainWindow::~MainWindow()
{
    this->freeAIThread();
    this->saveHistory();
    delete this->historyDialog;
    delete ui;
}

void MainWindow::loadComboBoxModelPicker(int index) {
    ui->comboBoxModelPicker->clear();
    if (index == 1) {
        QStringList serverUrlList;
        this->loadServerConfig(serverUrlList);
        if (serverUrlList.length()) {
            ui->comboBoxModelPicker->addItems(serverUrlList);
            ui->groupBoxChatter->setTitle("Chatter - Server");
        } else {
            ui->groupBoxChatter->setTitle("Chatter - NULL");
        }
        return;
    }
    
    QStringList modelList;
    if (!this->getLocalModelList(modelList)) {
        QMessageBox::warning(this, tr("Local model not found"), tr("Failed to find local model!"));
        return;
    }
    if (modelList.length()) {
        ui->comboBoxModelPicker->addItems(modelList);
        ui->groupBoxChatter->setTitle(tr("Chatter - ") + modelList.at(0));
    }
}

void MainWindow::adjustPushButtonSend() {
    if (ui->pushButtonLocker->text() == tr("Lock")) {
        if (ui->pushButtonSend->isEnabled())
            ui->pushButtonSend->setEnabled(false);
        return;
    }
    if (ui->pushButtonSend->isEnabled()) {
        if (this->model != nullptr
            || ui->textEditInputArea->toPlainText().trimmed().length() == 0)
            ui->pushButtonSend->setEnabled(false);
    } else {
        if (this->model == nullptr
            && ui->textEditInputArea->toPlainText().trimmed().length() != 0)
            ui->pushButtonSend->setEnabled(true);
    }
}

void MainWindow::saveHistory() {
    int cnt = this->answerList.length();
    if (cnt == 0)
        return;
    ChatHistory history;
    history.record(this->questionList, this->answerList, true);
    if (!QFileInfo(this->historyDir).isDir()) {
        QDir dir;
        dir.mkpath(this->historyDir);
    }
    if (!history.save(this->historyDir, this->dateStr)) {
        QMessageBox::critical(nullptr, tr("Auto save chat history"), tr("Failed to save chat history!"));
    }
}

void MainWindow::loadHistory(const QString& historyPath) {
    ChatHistory chatHistory;
    if (!chatHistory.load(historyPath))
        return;
    this->clear(true);
    this->dateStr = chatHistory.getDateStr();
    int cnt = chatHistory.getRecordsNumber();
    QString question;
    QString answer;
    for (int i = 0; i < cnt; ++ i) {
        chatHistory.read(i, question, answer);
        this->record(question, answer, false);
    }
}

void MainWindow::clear(bool saveHistory)
{
    if (saveHistory)
        this->saveHistory();
    ui->listWidgetQuestions->clear();
    this->questionList.clear();
    this->answerList.clear();
    ui->textBrowserOutputArea->clear();
    ui->textEditInputArea->clear();
    this->dateStr.clear();
    
    ui->pushButtonDisplayNewestChat->setEnabled(false);
}

bool MainWindow::isClean() {
    return ui->textBrowserOutputArea->toPlainText().length() == 0 &&
        ui->textEditInputArea->toPlainText().length() == 0 &&
        ui->listWidgetQuestions->count() == 0;
}

void MainWindow::record(const QString& question, const QString& answer, bool saveHistory) {
    int idx = question.indexOf('\n');
    if (idx < 0)
        ui->listWidgetQuestions->addItem(question);
    else
        ui->listWidgetQuestions->addItem(question.left(idx) + " ...");
    this->questionList.push_back(question);
    this->answerList.push_back(answer);
    if (saveHistory)
        this->saveHistory();
}

void MainWindow::readListItem(int index) {
    if (index < 0 || index >= ui->listWidgetQuestions->count())
        return;
    if (model != nullptr)
        ui->pushButtonDisplayNewestChat->setEnabled(true);
    ui->listWidgetQuestions->setCurrentRow(index);
    ui->textBrowserOutputArea->setText(
        QString("<user>\n%1\n</user>\n<ollama>\n%2\n</ollama>")
            .arg(this->questionList[index])
            .arg(this->answerList.at(index))
    );
}

bool MainWindow::getLocalModelList(QStringList& localModelList) {
    if (system("where ollama") != 0)
        return false;
    QProcess process;
    process.start("ollama", QStringList() << "list");
    if (!process.waitForFinished())
        return false;
    QStringList lines = QString(process.readAllStandardOutput()).split("\n");
    process.close();
    for (int i = 1; i < lines.length(); ++ i) {
        QString line = lines.at(i).trimmed();
        if (line.length() == 0)
            continue;
        localModelList.push_back(line.left(line.indexOf(" ")));
    }
    return true;
}

void MainWindow::askQuestion(const QString& question) {
    this->aiThread = new AIThread(this);
    this->model = new Model();
    QString jsonTemplate = "", keys = "";
    int currentRow = ui->listWidgetQuestions->currentRow();
    HistoryStorage historyStorage;
    int row = std::max(0, currentRow - historyStorage.getMaxSize() + 1);
    while (row <= currentRow) {
        historyStorage.push(this->questionList[row], this->answerList[row]);
        ++ row;
    }
    if (ui->comboBoxModelTypePicker->currentIndex() == 1) {
        int index = ui->comboBoxModelPicker->currentIndex();
        jsonTemplate = this->jsonList[index];
        keys = this->keysList[index];
    }
    this->model->setModelInfo(
        ui->comboBoxModelPicker->currentText(), jsonTemplate, keys
    );
    this->model->setQuestion(question, historyStorage.toString());
    this->model->moveToThread(this->aiThread);
    connect(this->aiThread, &QThread::started, this->model, &Model::work);
    connect(this->model, &Model::signal_updateOutputArea, this, &MainWindow::slot_updateOutputArea);
    connect(this->model, &Model::signal_finished, this, &MainWindow::slot_record);
    this->aiThread->start();
}

void MainWindow::freeAIThread() {
    if (this->aiThread != nullptr) {
        if (!this->aiThread->isInterruptionRequested())
            this->aiThread->requestInterruption();
        this->aiThread->quit();
        this->aiThread->wait();
        this->model->deleteLater();
        delete this->aiThread;
        this->aiThread = nullptr;
        this->model = nullptr;
    }
}

void MainWindow::loadServerConfig(QStringList& serverUrlList) {
    serverUrlList.clear();
    this->jsonList.clear();
    this->keysList.clear();
    QStringList configPaths;
    XIO::getPaths(ServersSettingsDialog::SERVERS_CONFIG_DIR, configPaths, false);
    QString serverUrl, json, keys;
    for (const QString& path : configPaths) {
        if (!ServersSettingsDialog::parse(path, serverUrl, json, keys))
            continue;
        serverUrlList.push_back(serverUrl);
        this->jsonList.push_back(json);
        this->keysList.push_back(keys);
    }
}

void MainWindow::on_actionNewChat_triggered()
{
    if (this->isClean())
        return;
    int choice = QMessageBox::question(nullptr, tr("Question"), tr("Start a new chat?"));
    if (choice == QMessageBox::Yes)
        this->clear(true);
}

void MainWindow::on_actionHistory_triggered()
{
    if (this->historyDialog == nullptr) {
        this->historyDialog = new HistoryDialog(this);
        connect(this->historyDialog, &HistoryDialog::signal_showHistory, 
                this, &MainWindow::slot_showHistory);
    }
    this->historyDialog->syncHistoryDirPath(this->historyDir);
    this->historyDialog->show();
}

void MainWindow::on_actionQuit_triggered()
{
    this->close();
}

void MainWindow::on_pushButtonSend_clicked()
{
    ui->pushButtonSend->setEnabled(false);
    ui->pushButtonStop->setEnabled(true);
    ui->pushButtonLocker->setEnabled(false);
    ui->pushButtonDisplayNewestChat->setEnabled(false);
    QString question = ui->textEditInputArea->toPlainText().trimmed();
    ui->textEditInputArea->clear();
    this->askQuestion(question);
}

void MainWindow::on_listWidgetQuestions_currentRowChanged(int currentRow)
{
    this->readListItem(currentRow);
}

void MainWindow::slot_showHistory(const ChatHistory& chatHistory, int index) {
    if (this->model != nullptr) {
        QMessageBox::critical(
            nullptr, tr("Invalid operation"), tr("Cannot read history while AI is still running!")
        );
        return;
    }
    this->clear(true);
    this->dateStr = chatHistory.getDateStr();
    int cnt = chatHistory.getRecordsNumber();
    QString question, answer;
    for (int i = 0; i < cnt; ++ i) {
        chatHistory.read(i, question, answer);
        this->record(question, answer, false);
    }
    if (index == -1)
        index = ui->listWidgetQuestions->count() - 1;
    this->readListItem(index);
}

void MainWindow::slot_updateOutputArea(const QString& originalQuestion, const QString& answer) {
    if (ui->pushButtonDisplayNewestChat->isEnabled() == false) {
        QScrollBar* scrollbar = ui->textBrowserOutputArea->verticalScrollBar();
        int currentPos = scrollbar->value();
        ui->textBrowserOutputArea->setText(
            QString("<user>\n%1\n</user>\n<ollama>\n%2\n")
                .arg(originalQuestion).arg(answer)
        );
        scrollbar->setValue(currentPos);
    }
}

void MainWindow::slot_record(const QString& question, const QString& answer) {
    this->record(question, answer, true);

    ui->listWidgetQuestions->setCurrentRow(ui->listWidgetQuestions->count() - 1);
    
    ui->pushButtonSend->setEnabled(true);
    ui->pushButtonStop->setEnabled(false);
    ui->pushButtonLocker->setEnabled(true);
    ui->pushButtonDisplayNewestChat->setEnabled(false);
    
    this->freeAIThread();
}

void MainWindow::slot_updateServersSettings() {
    if (ui->comboBoxModelTypePicker->currentIndex() == 1)
        this->loadComboBoxModelPicker(1);
}

void MainWindow::on_pushButtonLocker_clicked()
{
    QString currentText = ui->pushButtonLocker->text();
    if (currentText == tr("Unlock")) {
        ui->pushButtonLocker->setText(tr("Lock"));
        
        ui->comboBoxModelPicker->setEnabled(true);
        ui->comboBoxModelTypePicker->setEnabled(true);
        this->adjustPushButtonSend();
        
        return;
    }
    
    bool isURL = (ui->comboBoxModelTypePicker->currentText() == tr("Server URL"));
    QString modelNameOrServerURL = ui->comboBoxModelPicker->currentText();
    
    if (isURL) {
        ui->groupBoxChatter->setTitle("Chatter - Server");
    } else {
        if (ui->comboBoxModelPicker->findText(modelNameOrServerURL) < 0) {
            QMessageBox::warning(
                this, tr("Invalid model"), tr("Cannot find model \"%1\"!").arg(modelNameOrServerURL)
            );
            return;
        }
        ui->groupBoxChatter->setTitle(tr("Chatter - %1").arg(modelNameOrServerURL));
    }
    
    ui->pushButtonLocker->setText(tr("Unlock"));
    this->adjustPushButtonSend();
    ui->comboBoxModelPicker->setEnabled(false);
    ui->comboBoxModelTypePicker->setEnabled(false);
}

void MainWindow::on_pushButtonStop_clicked()
{
    this->freeAIThread();
}

void MainWindow::on_pushButtonDisplayNewestChat_clicked()
{
    ui->textBrowserOutputArea->clear();
    ui->pushButtonDisplayNewestChat->setEnabled(false);
    if (this->model == nullptr)
        this->readListItem(this->questionList.count() - 1);
}

void MainWindow::on_comboBoxModelTypePicker_currentIndexChanged(int index)
{
    this->loadComboBoxModelPicker(index);
}
void MainWindow::on_actionServersSettings_triggered()
{
    if (this->serversSettingsDialog == nullptr) {
        this->serversSettingsDialog = new ServersSettingsDialog(this);
        connect(this->serversSettingsDialog, &ServersSettingsDialog::signal_update,
                this, &MainWindow::slot_updateServersSettings);
    }
    this->serversSettingsDialog->show();
}

void MainWindow::on_textEditInputArea_textChanged()
{
    this->adjustPushButtonSend();
}
